package org.example.demo1.controller;

import org.example.demo1.ai.ChatRedisMemory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.awt.*;
import java.lang.annotation.Documented;
import java.security.PublicKey;
import java.util.List;
import java.util.Map;

@CrossOrigin
@RestController
public class SimpleAiController {

    @Autowired
    VectorStore vectorStore;
    //负责处理OpenAi
    private final ChatClient chatClient;

    @Autowired
    ChatRedisMemory chatRedisMemory;

    //对话记忆
    private final InMemoryChatMemory inMemoryChatMemory;

    public SimpleAiController(ChatClient chatClient, InMemoryChatMemory inMemoryChatMemory) {
        this.chatClient = chatClient;
        this.inMemoryChatMemory = inMemoryChatMemory;
    }

    //根据消息直接输出
    @PostMapping("/ai/call")
    public String call(@RequestBody Map<String, String> map) {
        String message = map.get("message");
        //此处相当于SpringAi封装了和大模型的交互 prompt开启一次聊天请求 call发请求 content返回响应
        return chatClient.prompt().user(message).call().content().trim();
    }

    //流式异步响应（类似于gpt边回答边出字 而不是一次全出完）
    @PostMapping(value = "/ai/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> streamChat(@RequestBody Map<String, String> map) {
        String message = map.get("message");
        return chatClient.prompt(message)
                .stream().content().map(content -> ServerSentEvent.<String>builder(content).event("message").build())
                .concatWithValues(ServerSentEvent.<String>builder("").build())
                .onErrorResume(e -> Flux.just(ServerSentEvent.builder("Error: " + e.getMessage()).event("error").build()));
    }

    //记忆化
    @GetMapping(value = "/ai/streamresp", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> streamResp(@RequestParam(value = "message", defaultValue = "Hello!") String message) {
        Flux<ServerSentEvent<String>> serverSentEventFlux = chatClient.prompt(message)
                //advisor会将用户前面的问题一起组装到prompt中，实现上下文记忆（windowsize是10）
                .advisors(new MessageChatMemoryAdvisor(inMemoryChatMemory, "123", 10), new SimpleLoggerAdvisor())
                .stream().content().map(content -> ServerSentEvent.builder(content).event("message").build())
                //问题回答结速标识,以便前端消息展示处理
                .concatWithValues(ServerSentEvent.builder("").build())
                .onErrorResume(e -> Flux.just(ServerSentEvent.builder("Error: " + e.getMessage()).event("error").build()));
        return serverSentEventFlux;
    }

    //整合知识库和自己提问的问题 一起向ai提问
    @GetMapping("ai/vectorStoreChat")
    public Flux<String>ollamaApi(@RequestParam(value = "message")String message){
        //先构造一个对知识库的信息检索
        SearchRequest searchRequest=SearchRequest.builder()
                .query(message)
                .topK(1) // 只取最相似的1条
                .build();
        //从知识库检索信息
        List<Document> documents=vectorStore.similaritySearch(searchRequest);
        //将检索到的信息和用户的提问一起构建prompt，起到一个增强查询的作用
        String targetMessage=String.format("已知信息:%s\n 用户提问:%s\n",documents.get(0).getText(),message);
        return chatClient.prompt(targetMessage).stream().content();
    }


    //整合知识库+记忆化查询+历史询问本地存储
    @GetMapping(value = "ai/finalResp")
    public Flux<ServerSentEvent<String>>streamResponse(@RequestParam(value = "message", defaultValue ="Hello!") String message,@RequestParam String sessionId){

        Flux<ServerSentEvent<String>>serverSentEventFlux=chatClient.prompt().user(message)
                .advisors(new MessageChatMemoryAdvisor(chatRedisMemory,sessionId,10))
                .stream().content().map(content -> ServerSentEvent.builder(content).event("message").build())
                //问题回答结速标识,以便前端消息展示处理
                .concatWithValues(ServerSentEvent.builder("").build())
                .onErrorResume(e -> Flux.just(ServerSentEvent.builder("Error: " + e.getMessage()).event("error").build()));
        return serverSentEventFlux;
    }
}
